#include <string.h>
#include <stdio.h>
#include <stdbool.h>
#include <unistd.h>
#include <stdlib.h>

#include "lwip/init.h"
#include "lwip/pbuf.h"
#include "lwip/ip.h"
#include "lwip/tcp.h"
#include "lwip/tcpip.h"
#include "lwip/altcp.h"
#include "lwip/altcp_tcp.h"
#include "lwip/err.h"
#include "lwip/timeouts.h"
#include "lwip/netif.h"
#include "lwip/init.h"
#include "netif/tapif.h"

#include "bsl_log.h"
#include "crypt_errno.h"
#include "crypt_eal_init.h"
#include "crypt_eal_rand.h"
#include "hitls_cert_init.h"
#include "hitls_crypt_init.h"

#include "altcp_tls_hitls.h"

/* Test configuration structure */
typedef struct {
    bool is_client;      /* true for client, false for server */
    bool with_auth_macro; /* whether to use AUTH_MODE macro */
    int32_t with_key;       /* 0: no, 1: yes, 2: encrypted */
    int32_t with_cert;      /* 0: no, 1: yes, 2: invalid */
    int32_t with_crl;       /* 0: no, 1: valid, 2: revoked, 3: invalid */
} test_config_t;

/* Network configuration */
#define SERVER_IP     "192.168.100.1"
#define SERVER_PORT   8888

static struct netif g_netif;
static bool send_flag = false;

/* File reading helper function */
static char* read_file_content(const char* filename)
{
    FILE *file = fopen(filename, "r");
    if (!file) {
        printf("Failed to open file: %s\n", filename);
        return NULL;
    }

    fseek(file, 0, SEEK_END);
    long file_size = ftell(file);
    fseek(file, 0, SEEK_SET);

    char *content = malloc(file_size + 1);
    if (!content) {
        printf("Failed to allocate memory for file content\n");
        fclose(file);
        return NULL;
    }

    size_t read_size = fread(content, 1, file_size, file);
    content[read_size] = '\0';
    fclose(file);

    return content;
}

/* Logging functions */
void BinLogFixLenFunc(uint32_t logId, uint32_t logLevel, uint32_t logType,
    void *format, void *para1, void *para2, void *para3, void *para4)
{
    (void)logLevel;
    (void)logType;
    printf("logId:%u\t", logId);
    printf(format, para1, para2, para3, para4);
    printf("\n");
}

void BinLogVarLenFunc(uint32_t logId, uint32_t logLevel, uint32_t logType,
    void *format, void *para)
{
    (void)logLevel;
    (void)logType;
    printf("logId:%u\t", logId);
    printf(format, para);
    printf("\n");
}

void BinLogInit(void)
{
    BSL_LOG_BinLogFuncs func = {0};
    func.fixLenFunc = BinLogFixLenFunc;
    func.varLenFunc = BinLogVarLenFunc;
    BSL_LOG_RegBinLogFunc(&func);
}

/* Client callback functions */
static void tls_client_err(void *arg, err_t err)
{
    (void)arg;
    printf("TLS客户端错误: %d, %s\n", err, lwip_strerr(err));
}

static err_t tls_client_recv(void *arg, struct altcp_pcb *pcb, struct pbuf *p, err_t err)
{
    (void)arg;
    (void)err;

    if (p == NULL) {
        printf("TLS客户端连接关闭: %d\n", err);
        altcp_close(pcb);
        return ERR_OK;
    }

    printf("客户端收到 %d 字节: %.*s\n", p->len, (int)p->len, (char *)p->payload);
    pbuf_free(p);
    return ERR_OK;
}

static err_t tls_client_sent(void *arg, struct altcp_pcb *pcb, u16_t len)
{
    (void)arg;
    (void)pcb;
    printf("客户端发送 %d 字节\n", len);
    return ERR_OK;
}

static err_t tls_client_connected(void *arg, struct altcp_pcb *pcb, err_t err)
{
    (void)arg;

    if (err != ERR_OK) {
        printf("TLS连接失败: %d\n", err);
        return err;
    }

    printf("TLS连接建立成功!\n");

    /* Send test message */
    const char *data = "Hello, Server from unified test!";
    altcp_write(pcb, data, strlen(data), TCP_WRITE_FLAG_COPY);
    altcp_output(pcb);

    return ERR_OK;
}

/* Server callback functions */
static void tls_server_err(void *arg, err_t err)
{
    (void)arg;
    printf("TLS服务器错误: %d, %s\n", err, lwip_strerr(err));
}

static err_t tls_server_recv(void *arg, struct altcp_pcb *pcb, struct pbuf *p, err_t err)
{
    (void)arg;
    (void)err;

    if (p == NULL) {
        printf("TLS服务器连接关闭: %d\n", err);
        altcp_close(pcb);
        return ERR_OK;
    }

    printf("服务器收到 %d 字节: %.*s\n", p->len, (int)p->len, (char *)p->payload);

    /* Echo received data */
    altcp_write(pcb, p->payload, p->len, TCP_WRITE_FLAG_COPY);
    altcp_output(pcb);

    pbuf_free(p);
    return ERR_OK;
}

static err_t tls_server_sent(void *arg, struct altcp_pcb *pcb, u16_t len)
{
    (void)arg;
    (void)len;

    if (!send_flag) {
        const char *data = "Hello, Client from unified test!";
        uint16_t data_len = strlen(data);
        uint16_t max_len = altcp_sndbuf(pcb);
        
        if (data_len > max_len) {
            data_len = max_len;
        }
        
        err_t err = altcp_write(pcb, data, data_len, TCP_WRITE_FLAG_COPY);
        if (err != ERR_OK) {
            printf("服务器发送数据失败: %d\n", err);
            return err;
        }
        
        send_flag = true;
        printf("服务器发送 %d 字节\n", data_len);
    }

    return ERR_OK;
}

static err_t tls_server_accept(void *arg, struct altcp_pcb *pcb, err_t err)
{
    (void)arg;

    if (err != ERR_OK) {
        printf("TLS服务器接受连接失败: %d\n", err);
        return err;
    }

    printf("TLS服务器接受新连接\n");

    altcp_recv(pcb, tls_server_recv);
    altcp_sent(pcb, tls_server_sent);
    altcp_err(pcb, tls_server_err);

    return ERR_OK;
}

void free_data(void *data)
{
    if (data == NULL) {
        return;
    }
    free(data);
}

/* Initialize client with auth mode */
static err_t init_tls_client_with_auth(test_config_t *config, int32_t auth_mode)
{
    struct altcp_pcb *pcb;
    struct altcp_tls_config *tls_config;
    ip_addr_t server_ip;

    /* Read certificate files based on configuration */
    char *ca_cert = NULL;
    char *client_cert = NULL;
    char *client_key = NULL;
    char *crl_data = NULL;

    /* Always read CA cert for server verification */
    ca_cert = read_file_content("../testcode/files/ca_and_int.pem");
    if (!ca_cert) {
        printf("Failed to read CA certificate\n");
        return ERR_VAL;
    }

    /* Read client certificate if needed */
    if (config->with_cert > 0) {
        if (config->with_cert == 2) {
            /* Use invalid certificate for testing */
            client_cert = read_file_content("../testcode/files/client.pem");
        } else {
            client_cert = read_file_content("../testcode/files/client.crt");
        }
        if (!client_cert) {
            printf("Failed to read client certificate\n");
            free_data(ca_cert);
            return ERR_VAL;
        }
    }

    /* Read client private key if needed */
    if (config->with_key > 0) {
        if (config->with_key == 2) {
            client_key = read_file_content("../testcode/files/client.enc.key");
        } else {
            client_key = read_file_content("../testcode/files/client.key");
        }
        if (!client_key) {
            printf("Failed to read client private key\n");
            free_data(ca_cert);
            free_data(client_cert);
            return ERR_VAL;
        }
    }

    /* Read CRL if needed */
    if (config->with_crl > 0) {
        switch (config->with_crl) {
            case 1:
                crl_data = read_file_content("../testcode/files/empty.crl");
                break;
            case 2:
                crl_data = read_file_content("../testcode/files/empty_and_client.crl");
                break;
        }
        if (!crl_data) {
            printf("Failed to read CRL data\n");
            free_data(ca_cert);
            free_data(client_cert);
            free_data(client_key);
            return ERR_VAL;
        }
    }

    /* Create TLS client configuration */
    if (config->with_key == 2 && client_key) {
        /* Encrypted private key */
        const char *key_pass = "123456";
        tls_config = altcp_tls_create_config_client_2wayauth(
            ca_cert, strlen(ca_cert) + 1,
            client_key, strlen(client_key) + 1, 
            key_pass, strlen(key_pass),
            client_cert, client_cert ? strlen(client_cert) + 1 : 0
#ifdef LWIP_ALTCP_HITLS_CRL
            , crl_data, crl_data ? strlen(crl_data) + 1 : 0
#endif
#ifdef LWIP_ALTCP_HITLS_AUTH_MODE
            , auth_mode
#endif
        );
    } else {
        /* Unencrypted private key or no key */
        tls_config = altcp_tls_create_config_client_2wayauth(
            ca_cert, strlen(ca_cert) + 1,
            client_key, client_key ? strlen(client_key) + 1 : 0, 
            NULL, 0,
            client_cert, client_cert ? strlen(client_cert) + 1 : 0
#ifdef LWIP_ALTCP_HITLS_CRL
            , crl_data, crl_data ? strlen(crl_data) + 1 : 0
#endif
#ifdef LWIP_ALTCP_HITLS_AUTH_MODE
            , auth_mode
#endif
        );
    }

    /* Clean up file contents */
    free_data(ca_cert);
    free_data(client_cert);
    free_data(client_key);
    free_data(crl_data);

    if (tls_config == NULL) {
        printf("Failed to create TLS client config\n");
        return ERR_MEM;
    }

    pcb = altcp_tls_new(tls_config, IPADDR_TYPE_ANY);
    if (pcb == NULL) {
        printf("Failed to create TLS client pcb\n");
        altcp_tls_free_config(tls_config);
        return ERR_MEM;
    }

    altcp_recv(pcb, tls_client_recv);
    altcp_sent(pcb, tls_client_sent);
    altcp_err(pcb, tls_client_err);

    /* Connect to server */
    ipaddr_aton(SERVER_IP, &server_ip);
    err_t err = altcp_connect(pcb, &server_ip, SERVER_PORT, tls_client_connected);
    if (err != ERR_OK) {
        printf("Failed to connect: %d\n", err);
        altcp_close(pcb);
        altcp_tls_free_config(tls_config);
        return err;
    }

    printf("正在连接到 %s:%d...\n", ipaddr_ntoa(&server_ip), SERVER_PORT);
    return ERR_OK;
}

/* Initialize server */
static err_t init_tls_server(test_config_t *config)
{
    struct altcp_tls_config *tls_config = NULL;
    struct altcp_pcb *pcb = NULL;

    /* Read certificate files based on configuration */
    char *server_cert = NULL;
    char *server_key = NULL;

    /* Read server certificate if needed */
    if (config->with_cert > 0) {
        if (config->with_cert == 2) {
            /* Use invalid certificate for testing */
            server_cert = read_file_content("../testcode/files/server.pem");
        } else {
            server_cert = read_file_content("../testcode/files/server.crt");
        }
        if (!server_cert) {
            printf("Failed to read server certificate\n");
            return ERR_VAL;
        }
    }

    /* Read server private key if needed */
    if (config->with_key > 0) {
        if (config->with_key == 2) {
            server_key = read_file_content("../testcode/files/server.enc.key");
        } else {
            server_key = read_file_content("../testcode/files/server.key");
        }
        if (!server_key) {
            printf("Failed to read server private key\n");
            free_data(server_cert);
            return ERR_VAL;
        }
    }

    /* Create TLS server configuration */
    if (config->with_key == 2 && server_key) {
        /* Encrypted private key */
        const char *key_pass = "123456";
        tls_config = altcp_tls_create_config_server_privkey_cert(
            server_key, strlen(server_key) + 1,
            key_pass, strlen(key_pass),
            server_cert, server_cert ? strlen(server_cert) + 1 : 0
        );
    } else if (config->with_key == 1 && server_key) {
        /* Unencrypted private key */
        tls_config = altcp_tls_create_config_server_privkey_cert(
            server_key, strlen(server_key) + 1,
            NULL, 0,
            server_cert, server_cert ? strlen(server_cert) + 1 : 0
        );
    } else {
        tls_config = altcp_tls_create_config_server_privkey_cert(NULL, 0, NULL, 0, NULL, 0);
    }

    /* Clean up file contents */
    free_data(server_cert);
    free_data(server_key);

    if (tls_config == NULL) {
        printf("Failed to create TLS server config\n");
        return ERR_MEM;
    }

    /* Create TLS server PCB */
    pcb = altcp_tls_new(tls_config, IPADDR_TYPE_ANY);
    if (pcb == NULL) {
        printf("Failed to create TLS server pcb\n");
        altcp_tls_free_config(tls_config);
        return ERR_MEM;
    }

    printf("TLS服务器配置成功\n");

    err_t err = altcp_bind(pcb, IP_ADDR_ANY, SERVER_PORT);
    if (err != ERR_OK) {
        printf("Failed to bind port: %d\n", err);
        altcp_close(pcb);
        altcp_tls_free_config(tls_config);
        return err;
    }

    /* Start listening */
    pcb = altcp_listen(pcb);
    if (pcb == NULL) {
        printf("Failed to listen\n");
        altcp_tls_free_config(tls_config);
        return ERR_MEM;
    }

    printf("TLS服务器启动成功\n");

    /* Set callback functions */
    altcp_accept(pcb, tls_server_accept);

    printf("TLS server listening on port %d\n", SERVER_PORT);
    return ERR_OK;
}

/* Initialize network interface */
static err_t init_network_interface(void)
{
    ip4_addr_t ipaddr;
    ip4_addr_t netmask;
    ip4_addr_t gw;

    /* Initialize LWIP */
    tcpip_init(NULL, NULL);
    printf("LWIP初始化完成\n");

    IP4_ADDR(&ipaddr, 192, 168, 100, 2);
    IP4_ADDR(&netmask, 255, 255, 255, 0);
    IP4_ADDR(&gw, 192, 168, 100, 1);

    /* Configure network interface */
    if (netif_add(&g_netif, &ipaddr, &netmask, &gw, NULL, tapif_init, tcpip_input) == NULL) {
        printf("Failed to configure network interface\n");
        return ERR_IF;
    }

    /* Start network interface */
    netif_set_default(&g_netif);
    netif_set_up(&g_netif);
    netif_set_link_up(&g_netif);
    printf("网络接口启动完成\n");

    return ERR_OK;
}

/* Initialize cryptographic libraries */
static err_t init_crypto_libraries(void)
{
    BinLogInit();

    int32_t ret = CRYPT_EAL_Init(CRYPT_EAL_INIT_RAND | CRYPT_EAL_INIT_PROVIDER);
    if (ret != 0) {
        printf("Failed to initialize crypt eal: %d\n", ret);
        return ERR_VAL;
    }

    ret = CRYPT_EAL_ProviderRandInitCtx(NULL, CRYPT_RAND_SHA256, "provider=default", NULL, 0, NULL);
    if (ret != CRYPT_EAL_ERR_DRBG_REPEAT_INIT && ret != 0) {
        printf("Failed to initialize random number generator: %d\n", ret);
        CRYPT_EAL_Cleanup(CRYPT_EAL_INIT_RAND | CRYPT_EAL_INIT_PROVIDER);
        return ERR_VAL;
    }

    HITLS_CertMethodInit();
    HITLS_CryptMethodInit();

    return ERR_OK;
}

/* Print usage information */
static void print_usage(const char *program_name)
{
    printf("用法: %s <client|server> <with_auth_macro> <auth> <with_key> <with_cert> <with_crl>\n", program_name);
    printf("参数说明:\n");
    printf("  client|server:        指定运行模式\n");
    printf("  with_auth_macro:      altcp_hitls是否带AUTH_MODE宏 (0=否, 1=是)\n");
    printf("  auth:                 认证模式 (0=无, 1=必需) - 仅客户端模式有效\n");
    printf("  with_key:             本端是否传入私钥 (0=否, 1=是, 2=加密私钥)\n");
    printf("  with_cert:            本端是否传入证书 (0=否, 1=是, 2=错误的证书)\n");
    printf("  with_crl:             本端是否传入CRL (0=否, 1=证书有效, 2=证书被吊销, 3=证书无效) - 仅客户端模式有效\n");
    printf("\n示例:\n");
    printf("  %s server 0 0 0 0 0  # 服务器模式,带AUTH_MODE宏,使用未加密私钥,建链成功\n", program_name);
    printf("  %s server 0 1 0 0 0  # 服务器模式,带AUTH_MODE宏,使用加密私钥,建链成功\n", program_name);
    printf("  %s server 0 1 0 0 0  # 服务器模式,带AUTH_MODE宏,使用加密私钥,传入错误的密钥,建链失败\n", program_name);
    printf("  %s server 0 0 0 0 0  # 服务器模式,带AUTH_MODE宏,客户端证书校验失败,断链\n", program_name);
    printf("  %s server 0 0 0 0 0  # 服务器模式,带AUTH_MODE宏,客户端证书被吊销,断链\n", program_name);
    printf("  %s server 0 1 0 0 0  # 服务器模式,不带AUTH_MODE宏,使用加密私钥,传入错误的密钥,校验失败,不断链\n", program_name);
    printf("  %s server 0 0 0 0 0  # 服务器模式,不带AUTH_MODE宏,客户端证书校验失败,不断链\n", program_name);
    printf("  %s server 0 0 0 0 0  # 服务器模式,不带AUTH_MODE宏,客户端证书被吊销,不断链\n", program_name);
    printf("  %s server 0 0 0 0 0  # 服务器模式,不带AUTH_MODE宏,使用未加密私钥,建链成功\n", program_name);

    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=0,不校验server证书,client不传入证书,server传入证书,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=0,不校验server证书,client传入证书,server传入证书,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=0,不校验server证书,server不传入证书,建链失败\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=1,校验server证书,使用未加密私钥,server证书校验成功,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=1,校验server证书,使用加密私钥,server证书校验成功,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=1,校验server证书,server证书校验失败,断链\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,带AUTH_MODE宏,isAuth=1,校验server证书,server证书被吊销,断链\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,不带AUTH_MODE宏,校验失败不断链,server传入证书,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,不带AUTH_MODE宏,校验失败不断链,server不传入证书,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,不带AUTH_MODE宏,校验失败不断链,server证书校验失败,建链成功\n", program_name);
    printf("  %s client 1 0 0 0 0  # 客户端模式,不带AUTH_MODE宏,校验失败不断链,server证书被吊销,建链成功\n", program_name);
}


int main(int argc, char *argv[])
{
    test_config_t config = {0};
    int32_t auth_mode = 0;

    /* Parse command line arguments */
    if (argc != 7) {
        print_usage(argv[0]);
        return -1;
    }

    /* Parse mode */
    if (strcmp(argv[1], "client") == 0) {
        config.is_client = true;
    } else if (strcmp(argv[1], "server") == 0) {
        config.is_client = false;
    } else {
        printf("错误: 模式必须是 'client' 或 'server'\n");
        print_usage(argv[0]);
        return -1;
    }

    /* Parse with_auth_macro */
    config.with_auth_macro = (atoi(argv[2]) != 0);

    /* Parse auth mode (client only) */
    auth_mode = atoi(argv[3]);
    if (auth_mode < 0 || auth_mode > 1) {
        printf("错误: 认证模式必须是 0 或 1\n");
        print_usage(argv[0]);
        return -1;
    }

    /* Parse with_key */
    config.with_key = atoi(argv[4]);
    if (config.with_key < 0 || config.with_key > 2) {
        printf("错误: with_key必须是 0, 1 或 2\n");
        print_usage(argv[0]);
        return -1;
    }

    /* Parse with_cert */
    config.with_cert = atoi(argv[5]);
    if (config.with_cert < 0 || config.with_cert > 2) {
        printf("错误: with_cert必须是 0, 1 或 2\n");
        print_usage(argv[0]);
        return -1;
    }

    /* Parse with_crl */
    config.with_crl = atoi(argv[6]);
    if (config.with_crl < 0 || config.with_crl > 3) {
        printf("错误: with_crl必须是 0, 1, 2 或 3\n");
        print_usage(argv[0]);
        return -1;
    }

    printf("测试配置:\n");
    printf("  模式: %s\n", config.is_client ? "客户端" : "服务器");
    printf("  AUTH_MODE宏: %s\n", config.with_auth_macro ? "启用" : "禁用");
    if (config.is_client) {
        printf("  认证模式: %s\n", auth_mode ? "必需" : "无");
    }
    const char* key_desc[] = {"无", "未加密", "加密"};
    printf("  私钥: %s\n", key_desc[config.with_key]);
    const char* cert_desc[] = {"无", "正常", "无效"};
    printf("  证书: %s\n", cert_desc[config.with_cert]);
    if (config.is_client) {
        const char* crl_desc[] = {"无", "有效", "被吊销", "无效"};
        printf("  CRL: %s\n", crl_desc[config.with_crl]);
    }
    printf("\n");

    /* Initialize network interface */
    err_t err = init_network_interface();
    if (err != ERR_OK) {
        return -1;
    }

    /* Initialize cryptographic libraries */
    err = init_crypto_libraries();
    if (err != ERR_OK) {
        return -1;
    }

    /* Initialize TLS client or server */
    if (config.is_client) {
        /* Store auth_mode in a global or pass it through config */
        /* For now, we'll modify the client init to accept auth_mode */
        err = init_tls_client_with_auth(&config, auth_mode);
    } else {
        err = init_tls_server(&config);
    }

    if (err != ERR_OK) {
        printf("Failed to initialize TLS %s: %d\n", config.is_client ? "client" : "server", err);
        CRYPT_EAL_Cleanup(CRYPT_EAL_INIT_RAND | CRYPT_EAL_INIT_PROVIDER);
        CRYPT_EAL_RandDeinit();
        return -1;
    }

    printf("TLS %s 初始化完成\n", config.is_client ? "客户端" : "服务器");

    /* Main loop */
    while (1) {
        sys_check_timeouts();
        usleep(1000);
    }

    /* Cleanup */
    CRYPT_EAL_Cleanup(CRYPT_EAL_INIT_RAND | CRYPT_EAL_INIT_PROVIDER);
    CRYPT_EAL_RandDeinit();
    return 0;
}